#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AST-based Code Injector
使用抽象语法树(AST)精确定位代码注入位置，避免简单文本搜索带来的错误
"""

import ast
import os
import random
import subprocess
import json
from pathlib import Path
from typing import List, Dict, Optional, Tuple, Any
from dataclasses import dataclass

from core.builtin_config_module.builtin_config import BuiltInConfig
from core.utils_module.language_utils import detect_language_for_file
from core.utils_module.logger import ObfuscationLogger


@dataclass
class InjectionPoint:
    """注入点信息"""
    file_path: str
    language: str
    position: int  # 字符位置
    line_number: int
    context: str  # 上下文类型: 'class', 'method', 'function', 'module'
    scope_name: str  # 作用域名称
    indentation: int  # 缩进级别


class ASTCodeInjector:
    """基于AST的代码注入器 - 精确定位注入位置"""

    def __init__(self, config: BuiltInConfig, logger: ObfuscationLogger):
        """
        初始化AST代码注入器

        Args:
            config: 统一配置对象
            logger: 日志记录器
        """
        self.config = config
        self.logger = logger
        self.injected_count = 0
        self.injection_points: List[InjectionPoint] = []

    def inject_from_file(self, content_file: str,
                        target_dirs: Optional[List[str]] = None,
                        use_ast: bool = True) -> bool:
        """
        从文件读取代码片段并使用AST注入

        Args:
            content_file: 代码片段文件路径
            target_dirs: 目标目录列表
            use_ast: 是否使用AST分析（为False时退回到正则模式）

        Returns:
            是否成功
        """
        try:
            # 使用参数或配置
            content_file = content_file or self.config.advanced_features.injection_content_file
            target_dirs = target_dirs or self.config.advanced_features.injection_target_dirs

            if not content_file:
                raise ValueError("未提供代码片段文件路径")

            if not os.path.exists(content_file):
                self.logger.log_error(f"代码片段文件不存在: {content_file}")
                return False

            self.logger.log_operation("AST代码注入", f"从文件读取: {content_file}")

            # 加载代码片段
            snippets = self._load_snippets(content_file)

            if not snippets:
                self.logger.log_warning("没有找到可用的代码片段")
                return False

            self.logger.log_operation("片段加载", f"加载了 {len(snippets)} 个代码片段")

            # 查找目标文件
            target_files = self._find_target_files(target_dirs)

            if not target_files:
                self.logger.log_error("没有找到目标文件")
                return False

            self.logger.log_operation("目标文件", f"找到 {len(target_files)} 个目标文件")

            # 分析注入点
            if use_ast:
                self.logger.log_operation("分析模式", "使用AST分析注入点")
                self._analyze_injection_points(target_files)
            else:
                self.logger.log_operation("分析模式", "使用传统正则模式")

            # 执行注入
            self.injected_count = 0
            for file_path in target_files:
                if use_ast and self.injection_points:
                    if self._inject_with_ast(file_path, snippets):
                        self.injected_count += 1
                else:
                    if self._inject_with_regex(file_path, snippets):
                        self.injected_count += 1

            self.logger.log_operation(
                "AST代码注入完成",
                f"成功注入 {self.injected_count}/{len(target_files)} 个文件"
            )

            return True

        except Exception as e:
            self.logger.log_error(f"AST代码注入失败: {e}")
            return False

    def _load_snippets(self, content_file: str) -> List[str]:
        """加载代码片段"""
        try:
            with open(content_file, 'r', encoding='utf-8') as f:
                content = f.read()

            # 按分隔符或行分割
            if '---' in content:
                snippets = [s.strip() for s in content.split('---') if s.strip()]
            else:
                snippets = [line.strip() for line in content.split('\n') if line.strip()]

            return snippets

        except Exception as e:
            self.logger.log_error(f"加载代码片段失败: {e}")
            return []

    def _find_target_files(self, target_dirs: List[str]) -> List[str]:
        """查找目标文件"""
        target_files = []

        if not target_dirs:
            target_dirs = [self.config.project_path]

        white_list = self.config.advanced_features.injection_white_list_folders

        for target_dir in target_dirs:
            if not os.path.exists(target_dir):
                self.logger.log_warning(f"目标目录不存在: {target_dir}")
                continue

            if os.path.isfile(target_dir):
                if self._should_inject_file(target_dir, white_list):
                    target_files.append(target_dir)
                continue

            for root, dirs, files in os.walk(target_dir):
                # 跳过不需要的目录
                dirs[:] = [d for d in dirs if not d.startswith('.')
                          and d not in ['build', 'dist', '__pycache__',
                                       'node_modules', 'DummyCode']]

                for file in files:
                    file_path = os.path.join(root, file)
                    if self._should_inject_file(file_path, white_list):
                        target_files.append(file_path)

        return target_files

    def _should_inject_file(self, file_path: str, white_list: List[str]) -> bool:
        """判断文件是否应该注入"""
        language = detect_language_for_file(file_path)
        if not language:
            return False

        if white_list:
            for white_dir in white_list:
                if white_dir in file_path:
                    return True
            return False

        return True

    def _analyze_injection_points(self, target_files: List[str]) -> None:
        """使用AST分析所有文件的注入点"""
        self.injection_points.clear()

        for file_path in target_files:
            language = detect_language_for_file(file_path)

            if language == 'python':
                points = self._analyze_python_ast(file_path)
            elif language == 'java':
                points = self._analyze_java_ast(file_path)
            elif language == 'swift':
                points = self._analyze_swift_ast(file_path)
            elif language == 'kotlin':
                points = self._analyze_kotlin_ast(file_path)
            elif language == 'objc':
                points = self._analyze_objc_ast(file_path)
            else:
                # 不支持的语言跳过
                continue

            self.injection_points.extend(points)

        self.logger.log_operation(
            "注入点分析",
            f"找到 {len(self.injection_points)} 个潜在注入点"
        )

    def _analyze_python_ast(self, file_path: str) -> List[InjectionPoint]:
        """分析Python文件的AST"""
        points = []

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            tree = ast.parse(content)

            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    # 函数定义
                    if hasattr(node, 'lineno'):
                        # 找到函数体的第一条语句位置
                        if node.body:
                            first_stmt = node.body[0]
                            point = InjectionPoint(
                                file_path=file_path,
                                language='python',
                                position=0,  # 稍后计算
                                line_number=first_stmt.lineno,
                                context='function',
                                scope_name=node.name,
                                indentation=first_stmt.col_offset if hasattr(first_stmt, 'col_offset') else 4
                            )
                            points.append(point)

                elif isinstance(node, ast.ClassDef):
                    # 类定义
                    if hasattr(node, 'lineno') and node.body:
                        # 在__init__方法后注入
                        for item in node.body:
                            if isinstance(item, ast.FunctionDef) and item.name == '__init__':
                                if item.body:
                                    first_stmt = item.body[0]
                                    point = InjectionPoint(
                                        file_path=file_path,
                                        language='python',
                                        position=0,
                                        line_number=first_stmt.lineno,
                                        context='class',
                                        scope_name=f"{node.name}.__init__",
                                        indentation=first_stmt.col_offset if hasattr(first_stmt, 'col_offset') else 8
                                    )
                                    points.append(point)
                                break

        except Exception as e:
            self.logger.log_warning(f"分析Python AST失败 {file_path}: {e}")

        return points

    def _analyze_java_ast(self, file_path: str) -> List[InjectionPoint]:
        """分析Java文件的AST"""
        points = []

        try:
            # 尝试使用javalang库
            try:
                import javalang
            except ImportError:
                self.logger.log_warning("javalang库未安装，使用回退方案")
                return self._analyze_java_with_regex(file_path)

            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            tree = javalang.parse.parse(content)

            for path, node in tree:
                if isinstance(node, javalang.tree.MethodDeclaration):
                    # 方法声明
                    if node.body and not node.modifiers and 'abstract' not in node.modifiers:
                        # 计算行号
                        lines = content[:node.position.column].count('\n')
                        point = InjectionPoint(
                            file_path=file_path,
                            language='java',
                            position=0,
                            line_number=lines + 2,  # 方法体开始
                            context='method',
                            scope_name=node.name,
                            indentation=8
                        )
                        points.append(point)

        except Exception as e:
            self.logger.log_warning(f"分析Java AST失败 {file_path}: {e}")
            return self._analyze_java_with_regex(file_path)

        return points

    def _analyze_java_with_regex(self, file_path: str) -> List[InjectionPoint]:
        """使用正则表达式分析Java（回退方案）"""
        import re
        points = []

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            # 查找方法定义
            method_pattern = re.compile(
                r'^\s*(public|private|protected)?\s*(static)?\s*\w+\s+(\w+)\s*\([^)]*\)\s*\{'
            )

            for i, line in enumerate(lines):
                match = method_pattern.match(line)
                if match:
                    method_name = match.group(3)
                    if method_name not in ['main', 'toString', 'equals', 'hashCode']:
                        point = InjectionPoint(
                            file_path=file_path,
                            language='java',
                            position=0,
                            line_number=i + 2,  # 下一行
                            context='method',
                            scope_name=method_name,
                            indentation=8
                        )
                        points.append(point)

        except Exception as e:
            self.logger.log_warning(f"正则分析Java失败 {file_path}: {e}")

        return points

    def _analyze_swift_ast(self, file_path: str) -> List[InjectionPoint]:
        """分析Swift文件的AST"""
        points = []

        try:
            # 使用swiftc的语法检查功能
            result = subprocess.run(
                ['swiftc', '-dump-ast', file_path],
                capture_output=True,
                text=True,
                timeout=10
            )

            if result.returncode == 0:
                ast_output = result.stdout
                # 解析AST输出找到函数和类
                lines = ast_output.split('\n')

                for line in lines:
                    if 'func_decl' in line:
                        # 提取函数信息
                        import re
                        match = re.search(r'"(\w+)"', line)
                        if match:
                            func_name = match.group(1)
                            # 简化：在函数开始处注入
                            point = InjectionPoint(
                                file_path=file_path,
                                language='swift',
                                position=0,
                                line_number=0,  # 需要进一步计算
                                context='function',
                                scope_name=func_name,
                                indentation=4
                            )
                            points.append(point)

        except (subprocess.TimeoutExpired, FileNotFoundError):
            self.logger.log_warning(f"Swift AST分析失败，使用回退方案: {file_path}")
            return self._analyze_swift_with_regex(file_path)

        return points

    def _analyze_swift_with_regex(self, file_path: str) -> List[InjectionPoint]:
        """使用正则表达式分析Swift（回退方案）"""
        import re
        points = []

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            # 查找函数定义
            func_pattern = re.compile(r'^\s*func\s+(\w+)\s*\([^)]*\).*\{')

            for i, line in enumerate(lines):
                match = func_pattern.match(line)
                if match:
                    func_name = match.group(1)
                    point = InjectionPoint(
                        file_path=file_path,
                        language='swift',
                        position=0,
                        line_number=i + 2,
                        context='function',
                        scope_name=func_name,
                        indentation=4
                    )
                    points.append(point)

        except Exception as e:
            self.logger.log_warning(f"正则分析Swift失败 {file_path}: {e}")

        return points

    def _analyze_kotlin_ast(self, file_path: str) -> List[InjectionPoint]:
        """分析Kotlin文件的AST"""
        points = []

        try:
            # 尝试使用kotlinc分析
            result = subprocess.run(
                ['kotlinc', '-Xprint-irs', file_path],
                capture_output=True,
                text=True,
                timeout=10
            )

            if result.returncode == 0:
                # 解析输出
                output = result.stdout
                # 简化处理，查找函数定义
                import re
                func_pattern = re.compile(r'FUN\s+name:(\w+)')
                matches = func_pattern.findall(output)

                for func_name in matches:
                    point = InjectionPoint(
                        file_path=file_path,
                        language='kotlin',
                        position=0,
                        line_number=0,  # 需要进一步计算
                        context='function',
                        scope_name=func_name,
                        indentation=4
                    )
                    points.append(point)

        except (subprocess.TimeoutExpired, FileNotFoundError):
            self.logger.log_warning(f"Kotlin AST分析失败，使用回退方案: {file_path}")
            return self._analyze_kotlin_with_regex(file_path)

        return points

    def _analyze_kotlin_with_regex(self, file_path: str) -> List[InjectionPoint]:
        """使用正则表达式分析Kotlin（回退方案）"""
        import re
        points = []

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            # 查找函数定义
            func_pattern = re.compile(r'^\s*fun\s+(\w+)\s*\([^)]*\).*\{')

            for i, line in enumerate(lines):
                match = func_pattern.match(line)
                if match:
                    func_name = match.group(1)
                    point = InjectionPoint(
                        file_path=file_path,
                        language='kotlin',
                        position=0,
                        line_number=i + 2,
                        context='function',
                        scope_name=func_name,
                        indentation=4
                    )
                    points.append(point)

        except Exception as e:
            self.logger.log_warning(f"正则分析Kotlin失败 {file_path}: {e}")

        return points

    def _analyze_objc_ast(self, file_path: str) -> List[InjectionPoint]:
        """分析ObjC文件的AST"""
        points = []

        try:
            # 使用clang分析
            result = subprocess.run(
                ['clang', '-Xclang', '-ast-dump', '-fsyntax-only', file_path],
                capture_output=True,
                text=True,
                timeout=10
            )

            if result.returncode == 0:
                ast_output = result.stdout
                # 解析AST输出
                import re
                method_pattern = re.compile(r'ObjCMethodDecl.*?(?:-)(\w+)')
                matches = method_pattern.findall(ast_output)

                for method_name in matches:
                    point = InjectionPoint(
                        file_path=file_path,
                        language='objc',
                        position=0,
                        line_number=0,  # 需要进一步计算
                        context='method',
                        scope_name=method_name,
                        indentation=4
                    )
                    points.append(point)

        except (subprocess.TimeoutExpired, FileNotFoundError):
            self.logger.log_warning(f"ObjC AST分析失败，使用回退方案: {file_path}")
            return self._analyze_objc_with_regex(file_path)

        return points

    def _analyze_objc_with_regex(self, file_path: str) -> List[InjectionPoint]:
        """使用正则表达式分析ObjC（回退方案）"""
        import re
        points = []

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
                lines = content.split('\n')

            # 查找方法实现
            method_pattern = re.compile(r'^[-+]\s*\([^)]+\)\s*(\w+)')

            for i, line in enumerate(lines):
                match = method_pattern.match(line)
                if match:
                    method_name = match.group(1)
                    # 找到方法体开始
                    for j in range(i, min(i + 10, len(lines))):
                        if '{' in lines[j]:
                            point = InjectionPoint(
                                file_path=file_path,
                                language='objc',
                                position=0,
                                line_number=j + 2,
                                context='method',
                                scope_name=method_name,
                                indentation=4
                            )
                            points.append(point)
                            break

        except Exception as e:
            self.logger.log_warning(f"正则分析ObjC失败 {file_path}: {e}")

        return points

    def _inject_with_ast(self, file_path: str, snippets: List[str]) -> bool:
        """使用AST信息注入代码"""
        try:
            # 找到此文件的注入点
            file_points = [p for p in self.injection_points if p.file_path == file_path]

            if not file_points:
                return False

            # 读取文件内容
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            # 选择一个注入点
            point = random.choice(file_points)

            # 选择一个代码片段
            snippet = random.choice(snippets)

            # 格式化代码片段
            formatted_snippet = self._format_snippet_with_context(
                snippet, point.language, point.context, point.indentation
            )

            # 在指定行注入
            if 0 < point.line_number <= len(lines):
                lines.insert(point.line_number - 1, formatted_snippet)
            else:
                # 如果行号无效，附加到文件末尾
                lines.append(formatted_snippet)

            # 写回文件
            with open(file_path, 'w', encoding='utf-8') as f:
                f.writelines(lines)

            self.logger.log_file_processing(
                file_path,
                f"AST注入成功 (位置: {point.context}:{point.scope_name})"
            )
            return True

        except Exception as e:
            self.logger.log_warning(f"AST注入失败 {file_path}: {e}")
            return False

    def _inject_with_regex(self, file_path: str, snippets: List[str]) -> bool:
        """使用正则表达式注入代码（回退方案）"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            language = detect_language_for_file(file_path)
            snippet = random.choice(snippets)

            # 格式化代码片段
            formatted_snippet = self._format_snippet(snippet, language)

            # 查找注入位置
            injection_pos = self._find_injection_position(content, language)

            if injection_pos is None:
                self.logger.log_warning(f"无法找到注入位置: {file_path}")
                return False

            # 插入代码
            new_content = content[:injection_pos] + formatted_snippet + content[injection_pos:]

            # 写回文件
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(new_content)

            self.logger.log_file_processing(file_path, "正则注入成功")
            return True

        except Exception as e:
            self.logger.log_warning(f"正则注入失败 {file_path}: {e}")
            return False

    def _format_snippet_with_context(self, snippet: str, language: str,
                                    context: str, indentation: int) -> str:
        """根据上下文格式化代码片段"""
        indent = ' ' * indentation

        # 添加注释和适当的缩进
        comment_styles = {
            'python': '#',
            'java': '//',
            'swift': '//',
            'kotlin': '//',
            'objc': '//'
        }

        comment = comment_styles.get(language, '//')

        # 根据上下文调整格式
        if context in ['method', 'function']:
            formatted = f"{indent}{comment} AST Injected code - {context}\n"
            for line in snippet.split('\n'):
                if line.strip():
                    formatted += f"{indent}{line}\n"
        elif context == 'class':
            formatted = f"{indent}{comment} AST Injected code - class member\n"
            for line in snippet.split('\n'):
                if line.strip():
                    formatted += f"{indent}{line}\n"
        else:
            formatted = f"{comment} AST Injected code\n{snippet}\n"

        return formatted

    def _format_snippet(self, snippet: str, language: str) -> str:
        """格式化代码片段（简单版本）"""
        if language in ['swift', 'objc', 'java', 'kotlin']:
            return f"\n    // Injected code\n    {snippet}\n"
        elif language == 'python':
            return f"\n    # Injected code\n    {snippet}\n"
        return f"\n{snippet}\n"

    def _find_injection_position(self, content: str, language: str) -> Optional[int]:
        """查找注入位置（回退方案）"""
        if language in ['swift', 'java', 'kotlin']:
            # 查找最后一个 }
            last_brace = content.rfind('}')
            if last_brace != -1:
                return last_brace
        elif language == 'objc':
            # 查找 @end 之前
            end_pos = content.rfind('@end')
            if end_pos != -1:
                line_start = content.rfind('\n', 0, end_pos)
                if line_start != -1:
                    return line_start
                return end_pos

        # 默认返回文件末尾
        return len(content)

    def get_statistics(self) -> dict:
        """获取注入统计"""
        return {
            'injected_count': self.injected_count,
            'injection_points_found': len(self.injection_points),
            'ast_success_rate': f"{(self.injected_count / max(len(self.injection_points), 1)) * 100:.1f}%"
        }